"""LangGraph Agent实现"""
from typing import TypedDict, Annotated, Sequence
from pathlib import Path
from datetime import datetime, timedelta, timezone
from langchain_core.messages import BaseMessage, AIMessage, ToolMessage, SystemMessage
from langchain_core.tools import BaseTool
from langchain_openai import ChatOpenAI
from langgraph.graph import StateGraph, END
from langgraph.checkpoint.memory import InMemorySaver
from config.config import OPENAI_API_KEY, OPENAI_BASE_URL, DEFAULT_MODEL


class AgentState(TypedDict):
    """Agent状态"""
    messages: Annotated[Sequence[BaseMessage], lambda x, y: x + y]


def get_beijing_time_str() -> tuple[str, str]:
    """获取北京时间和星期"""
    beijing_offset = timezone(timedelta(hours=8))
    beijing_time = datetime.now(beijing_offset)
    time_str = beijing_time.strftime("%Y年%m月%d日 %H:%M:%S")
    weekday = ["周一", "周二", "周三", "周四", "周五", "周六", "周日"][beijing_time.weekday()]
    return time_str, weekday


def inject_time_info(base_prompt: str) -> str:
    """注入时间信息到提示词"""
    time_str, weekday = get_beijing_time_str()
    return f"""【当前北京时间】{time_str} {weekday}

{base_prompt}"""


def load_system_prompt() -> str:
    """从文件加载系统提示词"""
    # 获取当前文件所在目录的父目录（my-agent目录）
    current_dir = Path(__file__).parent
    prompt_file = current_dir.parent / "system_prompt.txt"
    
    try:
        with open(prompt_file, "r", encoding="utf-8") as f:
            base_prompt = f.read().strip()
    except FileNotFoundError:
        # 如果文件不存在，返回默认提示词
        base_prompt = """你是一个智能助手。在调用任何工具之前，你必须先向用户说明调用该工具的理由和目的。
例如，在调用计算器工具之前，你应该说："我需要计算这个数学表达式，让我使用计算器工具来执行计算。"
在调用数据查询工具之前，你应该说："我需要查询相关数据来回答你的问题，让我使用数据查询工具。"
请确保在每次工具调用前都清楚地说明调用理由。"""
    
    return inject_time_info(base_prompt)


def create_agent(
    tools: list[BaseTool],
    model_name: str = None,
    system_prompt: str = None,
    api_key: str = None,
    base_url: str = None
) -> StateGraph:
    """创建LangGraph Agent
    
    Args:
        tools: 工具列表
        model_name: 模型名称（为None时使用环境变量配置的DEFAULT_MODEL）
        system_prompt: 系统提示词模板（为None时从文件加载）
        api_key: API密钥（为None时使用环境变量）
        base_url: API基础URL（为None时使用环境变量）
        
    Returns:
        配置好的StateGraph
    """
    # 如果未提供model_name，使用环境变量中的DEFAULT_MODEL
    if model_name is None:
        model_name = DEFAULT_MODEL
    
    # 初始化模型，绑定工具
    # 构建 LLM 初始化参数
    llm_kwargs = {
        "model": model_name,
        "temperature": 0
    }
    
    # 优先使用传入的参数，否则使用环境变量
    if api_key:
        llm_kwargs["api_key"] = api_key
    elif OPENAI_API_KEY:
        llm_kwargs["api_key"] = OPENAI_API_KEY
    
    if base_url:
        llm_kwargs["base_url"] = base_url
    elif OPENAI_BASE_URL:
        llm_kwargs["base_url"] = OPENAI_BASE_URL
    
    # 创建 LLM 实例
    import logging
    logger = logging.getLogger(__name__)
    logger.info(f"创建LLM实例: model={llm_kwargs.get('model')}, base_url={llm_kwargs.get('base_url')}, tools={len(tools)}")
    
    llm = ChatOpenAI(**llm_kwargs)
    
    # 只有当工具列表不为空时才绑定工具
    if tools:
        llm_with_tools = llm.bind_tools(tools)
    else:
        llm_with_tools = llm
    
    # 创建工具字典便于查找
    tool_dict = {tool.name: tool for tool in tools}
    
    # 自定义工具节点：显示工具调用提示
    def tool_node(state: AgentState):
        """工具节点，带提示信息"""
        from tools.sse_helper import send_sse_event
        
        messages = state["messages"]
        last_message = messages[-1]
        
        # 提取工具调用
        if isinstance(last_message, AIMessage) and last_message.tool_calls:
            tool_messages = []
            
            for tool_call in last_message.tool_calls:
                tool_name = tool_call["name"]
                tool_args = tool_call["args"]
                tool_id = tool_call["id"]
                
                # 发送工具调用开始事件
                send_sse_event("tool_call", {
                    "type": "start",
                    "tool_name": tool_name,
                    "tool_args": tool_args,
                    "tool_id": tool_id
                })
                
                # 执行工具调用
                tool = tool_dict[tool_name]
                result = tool.invoke(tool_args)
                
                # 发送工具调用结果事件
                send_sse_event("tool_call", {
                    "type": "result",
                    "tool_name": tool_name,
                    "tool_id": tool_id,
                    "result": str(result)
                })
                
                # 创建工具消息
                tool_messages.append(
                    ToolMessage(
                        content=str(result),
                        tool_call_id=tool_id
                    )
                )
            
            return {"messages": tool_messages}
        
        return {"messages": []}
    
    # Agent节点：调用模型
    def agent_node(state: AgentState):
        """Agent节点"""
        messages = state["messages"]
        
        # 获取系统提示词
        if system_prompt:
            current_prompt = inject_time_info(system_prompt)
        else:
            # 从文件加载
            current_prompt = load_system_prompt()
        
        # 检查消息列表是否已有系统消息，如果没有则添加
        messages_with_system = list(messages)
        if not messages_with_system or not isinstance(messages_with_system[0], SystemMessage):
            messages_with_system.insert(0, SystemMessage(content=current_prompt))
        
        try:
            response_message = llm_with_tools.invoke(messages_with_system)
            return {"messages": [response_message]}
        except Exception as e:
            import logging
            logger = logging.getLogger(__name__)
            logger.error(f"LLM调用失败: {str(e)}")
            logger.error(f"错误类型: {type(e).__name__}")
            if hasattr(e, 'response'):
                logger.error(f"响应状态码: {e.response.status_code if hasattr(e.response, 'status_code') else 'N/A'}")
                logger.error(f"响应内容: {e.response.text if hasattr(e.response, 'text') else 'N/A'}")
            raise
    
    # 路由函数：决定下一步操作
    def should_continue(state: AgentState) -> str:
        """判断是否需要继续调用工具"""
        messages = state["messages"]
        last_message = messages[-1]
        
        # 如果最后一条消息是AI消息且有工具调用，则调用工具
        if isinstance(last_message, AIMessage) and last_message.tool_calls:
            return "tools"
        # 否则结束
        return END
    
    # 构建图
    workflow = StateGraph(AgentState)
    
    # 添加节点
    workflow.add_node("agent", agent_node)
    workflow.add_node("tools", tool_node)
    
    # 设置入口点
    workflow.set_entry_point("agent")
    
    # 添加条件边
    workflow.add_conditional_edges(
        "agent",
        should_continue,
        {
            "tools": "tools",
            END: END
        }
    )
    
    # 从工具节点回到agent节点
    workflow.add_edge("tools", "agent")
    
    # 编译图
    checkpointer = InMemorySaver()
    return workflow.compile(checkpointer=checkpointer)

